import torch
from torchvision import transforms
import torch.nn as nn
from PIL import Image


torch.set_grad_enabled(False)

# 主函数：加载模型并预测
def sex_main(model_path, cropped_img, icon_size=32):
    """
    入参：
      - model_path: 权重路径（如 'assets/icon_gender_model.pth'）
      - cropped_img: PIL.Image（你上游传的是 icon）
      - icon_size: 预处理尺寸（需与训练一致）
    返回：
      {'prediction': 'female_icon'|'male_icon'|'未知', 'confidence': float, 'probabilities': {...}}
    """
    model, device, transform = get_gender_runtime(model_path, icon_size)

    # 预测
    prediction, confidence, probs = predict_image(
        model=model,
        image=cropped_img,   # 这里明确是 PIL.Image
        transform=transform,
        device=device,
        icon_size=(icon_size, icon_size)
    )

    # 低置信度兜底
    if confidence < 0.8:
        prediction = '未知'

    return {
        'prediction': prediction,
        'confidence': confidence,
        'probabilities': {
            'female_icon': float(probs[0]),
            'male_icon': float(probs[1])
        }
    }


# 预测单张图像
def predict_image(model, image, transform, device='cpu', icon_size=(32, 32)):
    """预测单张图像的性别图标"""
    model.eval()

    # 加载图像
    try:
        icon = image #Image.open(image_path).convert('RGB')
        # icon = extract_icon(image, icon_size)
    except Exception as e:
        print(f"Error loading image: {e}")
        # 返回默认结果
        return "female_icon", 50.0, [0.5, 0.5]

    image_tensor = transform(icon).unsqueeze(0)
    image_tensor = image_tensor.to(device)

    # 预测
    with torch.no_grad():
        output = model(image_tensor)
        probs = torch.nn.functional.softmax(output, dim=1)
        _, pred = torch.max(output, 1)

    # 获取预测结果
    class_names = ['female_icon', 'male_icon']
    prediction = class_names[pred.item()]
    confidence = probs[0][pred.item()].item() * 100

    return prediction, confidence, probs[0].cpu().numpy()



# ====== 性别模型全局缓存（只加载一次）======
GENDER_MODEL = None
GENDER_DEVICE = None
GENDER_TRANSFORM = None
GENDER_ICON_SIZE = None

def get_gender_runtime(model_path, icon_size=32):
    """
    懒加载 + 缓存：返回 (model, device, transform)
    只在第一次或 icon_size 变化时加载/构建
    """
    global GENDER_MODEL, GENDER_DEVICE, GENDER_TRANSFORM, GENDER_ICON_SIZE

    if GENDER_MODEL is not None and GENDER_ICON_SIZE == icon_size:
        return GENDER_MODEL, GENDER_DEVICE, GENDER_TRANSFORM

    # 选择设备
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    # 构建模型并加载权重（只做一次）
    model = IconClassifier(num_classes=2).to(device)
    # 如果你使用的是较新torch且权重是安全导出的，可加 weights_only=True
    state = torch.load(model_path, map_location=device)
    model.load_state_dict(state)
    model.eval()

    # 预处理（与训练时保持一致）
    transform = transforms.Compose([
        transforms.Resize((icon_size, icon_size)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    # 更新缓存
    GENDER_MODEL = model
    GENDER_DEVICE = device
    GENDER_TRANSFORM = transform
    GENDER_ICON_SIZE = icon_size

    # 可选：预热一下（加速首帧）
    try:
        with torch.no_grad():
            dummy = torch.zeros(1, 3, icon_size, icon_size, device=device)
            _ = model(dummy)
    except Exception:
        pass

    return GENDER_MODEL, GENDER_DEVICE, GENDER_TRANSFORM



# 定义轻量级CNN模型
class IconClassifier(nn.Module):
    def __init__(self, num_classes=2):
        super(IconClassifier, self).__init__()
        # 输入尺寸: 3x32x32
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 16x16x16

            nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 8x8x32

            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),  # 4x4x64
        )

        # 全连接分类器
        self.classifier = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Linear(64 * 4 * 4, 128),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(128, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x


# 图标提取函数 - 从图像右上角裁剪图标区域
def extract_icon(image, icon_size=(32, 32), crop_percent=0.4):
    """
    从图像右上角裁剪图标区域并调整大小

    参数:
        image: PIL Image对象
        icon_size: 图标输出尺寸
        crop_percent: 从右上角裁剪的比例
    """
    width, height = image.size
    # 计算裁剪区域 (右上角)
    crop_x = int(width * (1 - crop_percent))
    crop_y = 0
    crop_w = int(width * crop_percent)
    crop_h = int(height * crop_percent)

    # 裁剪并调整大小
    icon = image.crop((crop_x, crop_y, crop_x + crop_w, crop_y + crop_h))
    icon = icon.resize(icon_size, Image.Resampling.LANCZOS)
    return icon